import pandas as pd
from gensim.models import Word2Vec

word_model = Word2Vec.load('../model/word_vec.model')
class_keywords_df = pd.read_csv('../data/keywords_class.csv')
keywords = class_keywords_df['keywords']
labels = []
vectors = []

words = word_model.wv.vocab.keys()
for i in range(11):
    word_list = class_keywords_df.ix[i,:].values[0]
    for word in word_list.split(' '):
        if word in words:
            labels.append(str(i+1)+ '-'+word)
            vectors.append(word_model[word])

def plot_with_labels(low_dim_embs, labels, filename):
    assert low_dim_embs.shape[0] >= len(labels), 'More labels than embeddings'
    plt.figure(figsize=(18, 18))  # in inches
    for i, label in enumerate(labels):
      x, y = low_dim_embs[i, :]
      plt.scatter(x, y)
      plt.annotate(label,
                 xy=(x, y),
                 xytext=(5, 2),
                 textcoords='offset points',
                 ha='right',
                 va='bottom')

    plt.savefig(filename)

try:
    # pylint: disable=g-import-not-at-top
    from sklearn.manifold import TSNE
    import matplotlib.pyplot as plt

    plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
    plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

    tsne = TSNE(perplexity=30, n_components=2, init='pca', n_iter=5000, method='exact')
    low_dim_embs = tsne.fit_transform(vectors)
    plot_with_labels(low_dim_embs, labels, '../model/word2vec.png')

except ImportError as ex:
    print('Please install sklearn, matplotlib, and scipy to show embeddings.')
    print(ex)